[SPARK-57851][SQL] Shuffle-free single-task execution for small queries by viirya · Pull Request #56928 · apache/spark

viirya · 2026-07-01T15:19:50Z

What changes were proposed in this pull request?

This adds a conservative optimizer rule MarkSingleTaskExecution that marks small single-partition scans, optionally with a shuffle-inducing operator on top (sort, aggregate, distinct, window, limit/offset, expand) or an in-memory LocalRelation, as candidates for single-task execution. Such a scan reports a SinglePartition output partitioning, allowing EnsureRequirements to elide the shuffle that would otherwise be inserted before the operator on top.

Details:

The rule runs as the last optimizer batch (so it sees the final plan shape) and marks eligible LogicalRelation/LocalRelation nodes with a TreeNodeTag.
FileSourceStrategy/SparkStrategies propagate the mark to FileSourceScanExec/LocalTableScanExec.
FileSourceScanExec additionally gates on file count and size thresholds using the generic ScanFileListing, reports SinglePartition, and coalesces its input RDD to a single partition as a correctness backstop when the estimate does not match the runtime partition count.
LocalTableScanExec reads its data in a single partition and reports SinglePartition.
ExpandExec forwards SinglePartition from its child, since Expand only replicates rows within a partition and never moves rows across partitions.

The feature is controlled by new internal configs under spark.sql.optimizer.singleTaskExecution.* and is disabled by default. Join is intentionally left out for now and can be added as a follow-up; union is already covered by the existing spark.sql.unionOutputPartitioning.

This is part of the SPIP umbrella SPARK-56978 (Faster queries in local laptop mode), covering the shuffle-free local execution for small queries category.

Why are the changes needed?

For small, low-latency queries the fixed cost of a shuffle (scheduling, serialization, network) dominates the total runtime. When the input is already a single small partition, the shuffle inserted before a sort/aggregate/window is unnecessary and can be removed to reduce latency, without affecting correctness.

Does this PR introduce any user-facing change?

No. The optimization is behind internal configs (spark.sql.optimizer.singleTaskExecution.*) and is disabled by default.

How was this patch tested?

New MarkSingleTaskExecutionSuite (14 tests) covering:

the marking decision for the supported plan shapes;
SinglePartition output with no shuffle in the final physical plan;
empty-scan correctness (a global aggregation over an empty scan still returns a single row);
disabled-flag negatives (master flag and per-operator sub-flags);
ineligibility of unsupported shapes (join) and subquery expressions;
the leaf-node parallelism override disabling the local-relation case.

SQLConfSuite passes as a config-wiring regression check.

Was this patch authored or co-authored using generative AI tooling?

Yes, using Claude Code.

dongjoon-hyun · 2026-07-02T04:56:20Z


+  override def outputPartitioning: Partitioning = {
+    if (useSingleTask) {
+      SinglePartition


When rows.isEmpty=true, we need to return 0 partition instead of SinglePartition, don't we?

Good catch — this was a real correctness issue. The advertised SinglePartition did not match the zero-partition emptyRDD: with the shuffle elided, a global aggregation over an empty marked relation ran zero tasks and returned no rows instead of the single row expected on empty input. Fixed by producing one empty partition when the scan is marked, matching how maybeCoalesceInputRDD handles the all-files-pruned case in FileSourceScanExec, and added a test that failed before the fix.

dongjoon-hyun · 2026-07-02T05:03:21Z

+      val numFiles = selectedPartitions.totalNumberOfFiles
+      val numBytes = selectedPartitions.totalFileSize
+      numFiles >= minNumFiles && numFiles <= maxNumFiles &&
+        numBytes >= minNumBytes && numBytes <= maxPartitionBytes


In isLocalRelationEligible, get(SQLConf.LEAF_NODE_DEFAULT_PARALLELISM) is considered in a way. Do we need to consider SQLConf.LEAF_NODE_DEFAULT_PARALLELISM here too?

Right, file scans should respect it too. Moved the check up to apply() so the whole rule is skipped when a leaf-node parallelism override is set.

dongjoon-hyun · 2026-07-02T05:07:15Z

+  // across partitions, so when the single-task optimization is enabled and the child produces a
+  // single partition, we can forward the `SinglePartition` property to avoid an unneeded shuffle.
+  override def outputPartitioning: Partitioning = {
+    if (conf.getConf(SQLConf.SINGLE_TASK_EXECUTION_EXPAND) &&


If we read this session configuration every time, AQE seems to make this logic unstably.

Is this applied for all ExpandExec (both marked and unmarked)?

Addressed both points: the decision is now made at optimization time — MarkSingleTaskExecution also tags the Expand in a marked plan, and the planner passes it into ExpandExec as a constructor field — so outputPartitioning no longer reads the session conf at execution time, and the forwarding only applies within marked plans.

dongjoon-hyun · 2026-07-02T05:12:51Z

+   */
+  private[spark] lazy val maybeCoalesceInputRDD: RDD[InternalRow] = {
+    if (useSingleTaskExecution && inputRDD.getNumPartitions > 1) {
+      inputRDD.coalesce(1)


Is this consistent with outputPartitioning code path? In outputPartitioning code, bucketedScan is handled before useSingleTaskExecution.

Good catch — the combination was inconsistent: a marked bucketed scan would advertise HashPartitioning while maybeCoalesceInputRDD coalesced it to one partition. useSingleTaskExecution now returns false for bucketed scans, which keeps both code paths consistent by construction. Added a test.

dongjoon-hyun · 2026-07-02T05:16:21Z

Shall we add markedForSingleTaskExecution here to be safe?

dongjoon-hyun · 2026-07-02T05:19:56Z

+  val markTag: TreeNodeTag[Boolean] = TreeNodeTag[Boolean]("__single_task_execution")
+
+  private def get[T](entry: org.apache.spark.internal.config.ConfigEntry[T]): T =
+    SQLConf.get.getConf(entry)


Since Rule inherits SQLConfHelper already, shall we simply use conf.getConf directly?

spark/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/rules/Rule.scala

Line 24 in 1f80fa0

abstract class Rule[TreeType <: TreeNode[_]] extends SQLConfHelper with Logging {

spark/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/SQLConfHelper.scala

Line 32 in 1f80fa0

def conf: SQLConf = SQLConf.get

Done, removed the helper.

peter-toth · 2026-07-02T13:37:34Z

+    // Shuffle-inducing operators, allowed only when the matching sub-flag is enabled.
+    case _: Aggregate if enabled.aggregation =>
+      plan.children.forall(isSupportedShape(_, enabled))
+    case _: Distinct if enabled.aggregation =>


This Distinct branch is effectively dead here: ReplaceDistinctWithAggregate runs in the base Optimizer well before this last SparkOptimizer batch, so a Distinct has already been rewritten into an Aggregate by the time this rule sees the plan (the "scan + distinct" test in fact exercises the Aggregate branch, not this one). Harmless as defensive code, but worth either dropping it or adding a one-line note that Distinct is normally gone by now, so a future reader doesn't assume it can still reach this match.

Good catch — removed. It is indeed unreachable: ReplaceDistinctWithAggregate is non-excludable, so a Distinct can never survive to this last batch. While verifying this I found SubqueryAlias is dead for the same reason (EliminateSubqueryAliases runs in the non-excludable FinishAnalysis batch), so I dropped both and left a short note in the match.

Cross-checking against the shape checks this rule was derived from also surfaced a missing guard: an Aggregate containing a user-defined aggregation (e.g. functions.udaf or a typed Aggregator) passed the shape check, since it is an expression rather than an operator. It is excluded now via a new USER_DEFINED_AGGREGATION tree pattern — defensively, because an optimization that collapses partial/final aggregates separated by no exchange would skip the user's merge step.

dongjoon-hyun · 2026-07-02T17:41:40Z

  with ImplicitCastInputTypes
  with UserDefinedExpression {

+  final override val nodePatterns: Seq[TreePattern] = Seq(USER_DEFINED_AGGREGATION)


Is HiveUDAFFunction marked correctly like his?

spark/sql/hive/src/main/scala/org/apache/spark/sql/hive/hiveUDFs.scala

Lines 329 to 338 in 17e7e14

private[hive] case class HiveUDAFFunction(

name: String,

funcWrapper: HiveFunctionWrapper,

children: Seq[Expression],

isUDAFBridgeRequired: Boolean = false,

mutableAggBufferOffset: Int = 0,

inputAggBufferOffset: Int = 0)

extends TypedImperativeAggregate[HiveUDAFBuffer]

with HiveInspectors

with UserDefinedExpression {

Also, could you double-check V2Aggregator , too?

spark/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/V2Aggregator.scala

Lines 27 to 32 in 17e7e14

case class V2Aggregator[BUF <: java.io.Serializable, OUT](

aggrFunc: V2AggregateFunction[BUF, OUT],

children: Seq[Expression],

mutableAggBufferOffset: Int = 0,

inputAggBufferOffset: Int = 0)

extends TypedImperativeAggregate[BUF] with ImplicitCastInputTypes {

Good catch, it was not. Tagged HiveUDAFFunction with USER_DEFINED_AGGREGATION now — it is a TypedImperativeAggregate whose merge runs the Hive GenericUDAFEvaluator, so it needs the same guard.

Also missing — tagged it too. Worth noting V2Aggregator does not mix in UserDefinedExpression, so I went through the full set of aggregate expressions that run user code rather than keying off that trait: ScalaUDAF, ScalaAggregator, the two TypedAggregateExpressions, HiveUDAFFunction, and now V2Aggregator are all tagged. PythonUDAF is already covered by the existing PYTHON_UDF pattern. Added a test that instantiates V2Aggregator and asserts it carries the pattern.

viirya · 2026-07-02T23:03:20Z

The pyspark-connect-old-client failure is unrelated to this PR. It is a mode() doctest in python/pyspark/sql/connect/functions/builtin.py where the 4.0 client expects mode() WITHIN GROUP (ORDER BY col DESC) but the current server emits ORDER BY col (no DESC) — a client/server skew from a recent mode() change on master (SPARK-57329). This PR touches no Python or Connect mode code, and the failure reproduces on rerun. All other checks pass.

This adds a conservative optimizer rule `MarkSingleTaskExecution` that marks small single-partition scans, optionally with a shuffle-inducing operator on top (sort, aggregate, distinct, window, limit/offset, expand) or an in-memory `LocalRelation`, as candidates for single-task execution. Such a scan reports a `SinglePartition` output partitioning, allowing `EnsureRequirements` to elide the shuffle that would otherwise be inserted before the operator on top. The rule runs as the last optimizer batch and marks eligible `LogicalRelation`/`LocalRelation` nodes with a `TreeNodeTag`. The planning strategies propagate the mark to `FileSourceScanExec`/`LocalTableScanExec`. `FileSourceScanExec` additionally gates on file count and size thresholds using the generic `ScanFileListing`, reports `SinglePartition`, and coalesces its input RDD to a single partition as a correctness backstop. `ExpandExec` forwards `SinglePartition` from its child, since Expand never moves rows across partitions. The feature is controlled by new internal configs under `spark.sql.optimizer.singleTaskExecution.*` and is disabled by default. Join is intentionally left out for now; union is already covered by the existing `spark.sql.unionOutputPartitioning`. This is part of the SPIP umbrella SPARK-56978 (Faster queries in local laptop mode), covering the shuffle-free local execution category. For small, low-latency queries the fixed cost of a shuffle (scheduling, serialization, network) dominates. When the input is already a single small partition, the shuffle before a sort/aggregate/window is unnecessary and can be removed to reduce latency. No. The optimization is behind internal configs and is disabled by default. New `MarkSingleTaskExecutionSuite` (14 tests) covering the marking decision, `SinglePartition` output with no shuffle, empty-scan correctness, disabled-flag negatives, join/subquery ineligibility, and the leaf-parallelism override. `SQLConfSuite` passes as a config-wiring regression check. Co-authored-by: Isaac

Co-authored-by: Claude Code

- LocalTableScanExec: produce one empty partition for an empty marked relation to match the advertised SinglePartition; a zero-partition RDD with the shuffle elided returns no rows for a global aggregation. - FileSourceScanLike: exclude bucketed scans from single-task execution; coalescing would invalidate their HashPartitioning. - ExpandExec: decide SinglePartition forwarding at planning time from the markTag instead of reading the session conf at execution time, and only within marked plans. - MarkSingleTaskExecution: skip the whole rule when a leaf-node parallelism override is set, so file scans respect it too; use the inherited conf.getConf instead of a private helper. - FileSourceScanExec.doCanonicalize: retain markedForSingleTaskExecution. Co-authored-by: Claude Code

The new useSingleTask constructor parameter appended ", false" to every Expand node's explain arguments, breaking the TPC-DS plan stability golden files for rollup queries. Show the flag only when it is set. Co-authored-by: Claude Code

…TaskExecution - Add the USER_DEFINED_AGGREGATION tree pattern, tag ScalaUDAF, ScalaAggregator and the typed aggregate expressions with it, and add it to the rule's unsupported patterns. This is a defensive guard: an optimization that collapses partial and final aggregates separated by no exchange would skip the user's merge step. - Drop the Distinct and SubqueryAlias cases from isSupportedShape: both are rewritten away by non-excludable rules (ReplaceDistinctWithAggregate, EliminateSubqueryAliases) before this last optimizer batch runs. Co-authored-by: Claude Code

Co-authored-by: Claude Code

These are the remaining user-defined aggregate expressions whose merge step runs user code, so plans containing them must also be excluded from single-task execution. PythonUDAF is already covered by the PYTHON_UDF pattern. Note V2Aggregator does not mix in UserDefinedExpression, so it would be missed by a trait-based scan. Co-authored-by: Claude Code

Co-authored-by: Claude Code

dongjoon-hyun

IIUC, the current design seems to let tags escape to other queries. For example,

Assume that t is cached.
A query SELECT col FROM t ORDER BY col marked t as __single_task_execution. The cached t is mutated.
A user disables this feature via singleTaskExecution.enabled=false.
A new query executes t JOIN huge_table. The cached and marked t can cause a problem with the single partition.

The above issue persists until refreshTable. Maybe, we need to redesign this like a physical plan like DisableUnnecessaryBucketedScan instead of logical plans.

viirya · 2026-07-07T01:08:30Z

Thanks, I dug into this. I couldn't reproduce the leak as described, but I agree the underlying design is fragile and worth changing.

What I found when I tried to reproduce it:

Marking is an in-place setTagValue on the LogicalRelation, and tags do propagate across clone() via copyTagsFrom — so your read of the mechanism is right.
But two implementation details happen to block the specific scenario. First, each sql("... FROM t") re-resolves a fresh LogicalRelation from the catalog, so a tag set by one query never reaches the next query's node (I confirmed a feature-ON ORDER BY query followed by a feature-OFF JOIN over the same table leaves the join's scans untagged). Second, when t is cached, the relation is already replaced by InMemoryRelation by the time this rule runs, so step 2 never actually tags the cached plan — I checked the cache-held logical plan (cacheBuilder.logicalPlan) after a markable feature-ON query and its LogicalRelation carries no tag.

So the leak doesn't happen today, but only because of that chain of coincidences, not by design — which is exactly your point about mutating logical plans. Rather than a full physical-plan redesign, would it be acceptable to make the marking copy-based so the rule never mutates a node it was handed? i.e. mark a lr.copy() instead of lr itself (.copy() preserves the exprIds, so it's semantics-preserving). That restores the "an optimizer rule doesn't mutate shared input nodes" invariant at essentially no cost and guards against this breaking if those coincidences ever change. Happy to go the DisableUnnecessaryBucketedScan-style physical route instead if you'd prefer the stronger guarantee.

dongjoon-hyun · 2026-07-07T20:03:12Z

Thank you for verifying that. Given that, the copy-based marking sounds good to me, @viirya .

Address the review concern that the rule tags LogicalRelation nodes in place: a tag set on a shared node propagates through TreeNode.clone via copyTagsFrom and could leak the marking into unrelated plans. The rule now clones the plan and tags the clone, never touching the nodes it was handed. Note that marking per-node copies does not work: a copy differing only in tags is structurally equal to the original, so withNewChildren's fastEquals short-circuit discards the copy and keeps the original (untagged) node. Cloning the whole plan avoids all rebuild paths. The clone only happens for plans that will actually be marked, which are small by construction. Adds a regression test asserting the input plan's nodes stay untagged. Co-authored-by: Claude Code

viirya · 2026-07-07T21:19:22Z

Implemented, with one twist worth noting. Marking a per-node copy didn't survive: a copy that differs only in tags is structurally equal to the original, so withNewChildren's fastEquals short-circuit discards the tagged copy and keeps the original untagged node. Instead the rule now clones the whole plan and tags the clone in place, which achieves the same goal — the rule never mutates a node it was handed — without fighting the tree-rebuilding APIs. The clone only happens for plans that will actually be marked, which are small by construction. Also added a regression test asserting the input plan's nodes stay untagged after the rule runs.

dongjoon-hyun

+1, LGTM.

viirya requested a review from dtenedor July 1, 2026 15:23

viirya force-pushed the single-node-execution branch from e2fff1a to 7e18780 Compare July 1, 2026 15:26

dongjoon-hyun reviewed Jul 2, 2026

View reviewed changes

peter-toth reviewed Jul 2, 2026

View reviewed changes

peter-toth approved these changes Jul 2, 2026

View reviewed changes

dongjoon-hyun reviewed Jul 2, 2026

View reviewed changes

dongjoon-hyun mentioned this pull request Jul 2, 2026

[SPARK-57885][SQL] Promote EmptyRDDWithPartitions to a top-level class #56967

Closed

viirya added 8 commits July 6, 2026 10:31

Set NOT_APPLICABLE binding policy on singleTaskExecution configs

e022de7

Co-authored-by: Claude Code

Fix scalastyle line length in MarkSingleTaskExecutionSuite

d44f731

Co-authored-by: Claude Code

Fix scalastyle import ordering in MarkSingleTaskExecutionSuite

f4cfaaa

Co-authored-by: Claude Code

viirya force-pushed the single-node-execution branch from 4de7031 to f4cfaaa Compare July 6, 2026 17:34

dongjoon-hyun reviewed Jul 6, 2026

View reviewed changes

dongjoon-hyun approved these changes Jul 7, 2026

View reviewed changes

	private[hive] case class HiveUDAFFunction(
	name: String,
	funcWrapper: HiveFunctionWrapper,
	children: Seq[Expression],
	isUDAFBridgeRequired: Boolean = false,
	mutableAggBufferOffset: Int = 0,
	inputAggBufferOffset: Int = 0)
	extends TypedImperativeAggregate[HiveUDAFBuffer]
	with HiveInspectors
	with UserDefinedExpression {

	case class V2Aggregator[BUF <: java.io.Serializable, OUT](
	aggrFunc: V2AggregateFunction[BUF, OUT],
	children: Seq[Expression],
	mutableAggBufferOffset: Int = 0,
	inputAggBufferOffset: Int = 0)
	extends TypedImperativeAggregate[BUF] with ImplicitCastInputTypes {

Uh oh!

Conversation

viirya commented Jul 1, 2026

What changes were proposed in this pull request?

Why are the changes needed?

Does this PR introduce any user-facing change?

How was this patch tested?

Was this patch authored or co-authored using generative AI tooling?

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

viirya commented Jul 2, 2026

Uh oh!

dongjoon-hyun left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

viirya commented Jul 7, 2026

Uh oh!

dongjoon-hyun commented Jul 7, 2026

Uh oh!

viirya commented Jul 7, 2026

Uh oh!

dongjoon-hyun left a comment

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

dongjoon-hyun left a comment •

edited

Loading